package com.zh.stargraph2.core.service.impl;

import com.alibaba.fastjson2.JSON;
import com.zh.stargraph2.core.pojo.ComfyuiTask;
import com.zh.stargraph2.core.service.RedisService;
import com.zh.stargraph2.core.service.WsNoticeService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Service;

import java.util.HashMap;

@Service
public class WsNoticeServiceImpl implements WsNoticeService {

    public final static String COMFYUI_QUEUE_TOPIC = "/topic/messages";

    @Autowired
    SimpMessagingTemplate simpMessagingTemplate;

    @Autowired
    RedisService redisService;

    @Override
    public void sendProgress(String promptId, HashMap<String, Object> data) {
        ComfyuiTask task = redisService.getStartedTask(promptId);
        if (task != null) {
            data.put("type", "progress");
            sendToUser(task.getWsClientId(), JSON.toJSONString(data));
        }
    }

    /**
     * 发送给对应用户
     * @param wsClientId
     * @param message
     */
    private void sendToUser(String wsClientId, String message) {
        simpMessagingTemplate.convertAndSendToUser(wsClientId, COMFYUI_QUEUE_TOPIC, message);
    }

    /**
     * 发送给所有用户
     * @param message
     */
    private void sendToAll(String message) {
        simpMessagingTemplate.convertAndSend(COMFYUI_QUEUE_TOPIC, message);
    }


}
